import os
import datetime
import ipaddress
from cryptography import x509
from cryptography.x509.oid import NameOID
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
from cryptography.hazmat.primitives.hashes import SHA256

# 配置参数
OUTPUT_DIR = "certificates"
KEY_SIZE = 256  # 密钥大小，可选256, 384, 521
CA_EXPIRY_YEARS = 10
CERT_EXPIRY_YEARS = 5

def generate_ec_key(key_size=256):
    """生成ECDSA私钥"""
    curve = {
        256: ec.SECP256R1(),
        384: ec.SECP384R1(),
        521: ec.SECP521R1()
    }.get(key_size, ec.SECP256R1())
    private_key = ec.generate_private_key(curve, default_backend())
    return private_key

def generate_ca_cert(private_key):
    """生成CA证书"""
    subject = issuer = x509.Name([
        x509.NameAttribute(NameOID.COUNTRY_NAME, "CN"),
        x509.NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, "Shanghai"),
        x509.NameAttribute(NameOID.LOCALITY_NAME, "Shanghai"),
        x509.NameAttribute(NameOID.ORGANIZATION_NAME, "My Device CA"),
        x509.NameAttribute(NameOID.COMMON_NAME, "device-ca.example.com"),
    ])
    
    cert = x509.CertificateBuilder().subject_name(
        subject
    ).issuer_name(
        issuer
    ).public_key(
        private_key.public_key()
    ).serial_number(
        x509.random_serial_number()
    ).not_valid_before(
        datetime.datetime.utcnow()
    ).not_valid_after(
        datetime.datetime.utcnow() + datetime.timedelta(days=365 * CA_EXPIRY_YEARS)
    ).add_extension(
        x509.BasicConstraints(ca=True, path_length=None), critical=True,
    ).add_extension(
        x509.KeyUsage(
            digital_signature=True,
            content_commitment=False,
            key_encipherment=False,
            data_encipherment=False,
            key_agreement=False,
            key_cert_sign=True,
            crl_sign=True,
            encipher_only=False,
            decipher_only=False
        ), critical=True
    ).add_extension(
        x509.SubjectKeyIdentifier.from_public_key(private_key.public_key()),
        critical=False
    ).add_extension(
        x509.AuthorityKeyIdentifier.from_issuer_public_key(private_key.public_key()),
        critical=False
    ).sign(private_key, hashes.SHA256(), default_backend())
    
    return cert

def generate_device_cert(private_key, ca_private_key, ca_cert, device_id):
    """生成设备证书"""
    subject = x509.Name([
        x509.NameAttribute(NameOID.COUNTRY_NAME, "CN"),
        x509.NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, "Shanghai"),
        x509.NameAttribute(NameOID.LOCALITY_NAME, "Shanghai"),
        x509.NameAttribute(NameOID.ORGANIZATION_NAME, "Device Manufacturer"),
        x509.NameAttribute(NameOID.COMMON_NAME, f"device-{device_id}.example.com"),
        x509.NameAttribute(NameOID.SERIAL_NUMBER, device_id),
    ])
    
    cert = x509.CertificateBuilder().subject_name(
        subject
    ).issuer_name(
        ca_cert.subject
    ).public_key(
        private_key.public_key()
    ).serial_number(
        x509.random_serial_number()
    ).not_valid_before(
        datetime.datetime.utcnow()
    ).not_valid_after(
        datetime.datetime.utcnow() + datetime.timedelta(days=365 * CERT_EXPIRY_YEARS)
    ).add_extension(
        x509.BasicConstraints(ca=False, path_length=None), critical=True,
    ).add_extension(
        x509.KeyUsage(
            digital_signature=True,
            content_commitment=False,
            key_encipherment=True,
            data_encipherment=False,
            key_agreement=False,
            key_cert_sign=False,
            crl_sign=False,
            encipher_only=False,
            decipher_only=False
        ), critical=True
    ).add_extension(
        x509.ExtendedKeyUsage([
            x509.OID_CLIENT_AUTH,
            x509.OID_SERVER_AUTH,
        ]), critical=False
    ).add_extension(
        x509.SubjectAlternativeName([
            x509.DNSName(f"device-{device_id}.example.com"),
            x509.DNSName("localhost"),
            x509.IPAddress(ipaddress.IPv4Address("127.0.0.1")),  # 修复类型错误
        ]), critical=False
    ).add_extension(
        x509.AuthorityKeyIdentifier.from_issuer_public_key(ca_private_key.public_key()),
        critical=False
    ).sign(ca_private_key, hashes.SHA256(), default_backend())
    
    return cert

def generate_shared_secret(private_key, peer_public_key):
    """生成共享密钥"""
    shared_key = private_key.exchange(ec.ECDH(), peer_public_key)
    # 使用HKDF派生密钥（可选，确保密钥长度一致）
    derived_key = HKDF(
        algorithm=SHA256(),
        length=32,
        salt=None,
        info=b"shared secret",
        backend=default_backend()
    ).derive(shared_key)
    return derived_key

def format_shared_secret_as_c_array(shared_secret, var_name):
    """将共享密钥格式化为C语言数组"""
    hex_values = ", ".join(f"0x{byte:02x}" for byte in shared_secret)
    return f"// {var_name}\nconst uint8_t {var_name}[] = {{ {hex_values} }};\n"

def save_to_file(data, path, is_private_key=False, is_public_key=False):
    """保存数据到文件"""
    os.makedirs(os.path.dirname(path), exist_ok=True)
    with open(path, "wb") as f:
        if is_private_key:
            f.write(data.private_bytes(
                encoding=serialization.Encoding.PEM,
                format=serialization.PrivateFormat.PKCS8,
                encryption_algorithm=serialization.NoEncryption()
            ))
        elif is_public_key:
            f.write(data.public_bytes(
                encoding=serialization.Encoding.PEM,
                format=serialization.PublicFormat.SubjectPublicKeyInfo
            ))
        elif isinstance(data, x509.Certificate):
            f.write(data.public_bytes(serialization.Encoding.PEM))
        else:
            f.write(data)
    print(f"已保存: {path}")

def format_cert_as_c_variable(cert, var_name):
    """将证书格式化为C语言变量，并在上方添加原始证书注释"""
    cert_pem = cert.public_bytes(serialization.Encoding.PEM).decode("utf-8")
    cert_lines = cert_pem.splitlines()
    formatted_cert = f'/**\n{cert_pem}\n*/\n'
    formatted_cert += f'// {var_name}\nconst char *{var_name} =\n'
    for line in cert_lines:
        formatted_cert += f'    "{line}\\r\\n"\n'
    formatted_cert += '    ;\n'
    return formatted_cert

def format_key_as_c_variable(key, var_name, is_private_key=False):
    """将密钥格式化为C语言变量，并在上方添加原始密钥注释"""
    if is_private_key:
        key_pem = key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.PKCS8,
            encryption_algorithm=serialization.NoEncryption()
        ).decode("utf-8")
    else:
        key_pem = key.public_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PublicFormat.SubjectPublicKeyInfo
        ).decode("utf-8")
    
    key_lines = key_pem.splitlines()
    formatted_key = f'/**\n{key_pem}\n*/\n'
    formatted_key += f'// {var_name}\nconst char *{var_name} =\n'
    for line in key_lines:
        formatted_key += f'    "{line}\\r\\n"\n'
    formatted_key += '    ;\n'
    return formatted_key

def append_to_log(log_path, content):
    """将内容追加到日志文件"""
    with open(log_path, "a", encoding="utf-8") as log_file:
        log_file.write(content + "\n")

def main():
    """主函数"""
    print("开始生成设备CA证书和示例设备证书...")
    
    # 创建输出目录
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    log_path = os.path.join(OUTPUT_DIR, "log.txt")
    if os.path.exists(log_path):
        os.remove(log_path)  # 清空旧日志文件
    
    # 生成设备CA
    print("\n生成设备CA...")
    ca_private_key = generate_ec_key(KEY_SIZE)
    ca_cert = generate_ca_cert(ca_private_key)
    
    # 保存设备CA
    save_to_file(ca_private_key, f"{OUTPUT_DIR}/ca/ca_private_key.pem", is_private_key=True)
    save_to_file(ca_cert, f"{OUTPUT_DIR}/ca/ca_cert.pem")
    save_to_file(ca_private_key.public_key(), f"{OUTPUT_DIR}/ca/ca_public_key.pem", is_public_key=True)
    
    # 打印CA证书和密钥为C语言格式并写入日志
    ca_cert_c = format_cert_as_c_variable(ca_cert, "ca_certificates")
    ca_private_key_c = format_key_as_c_variable(ca_private_key, "ca_private_key", is_private_key=True)
    ca_public_key_c = format_key_as_c_variable(ca_private_key.public_key(), "ca_public_key")
    append_to_log(log_path, ca_cert_c)
    append_to_log(log_path, ca_private_key_c)
    append_to_log(log_path, ca_public_key_c)
    
    # 生成并保存示例设备证书
    print("\n生成示例设备证书...")
    device_keys = {}  # 存储设备的私钥和公钥
    for device_id in ["001", "002"]:
        print(f"\n生成设备 {device_id} 的证书...")
        device_private_key = generate_ec_key(KEY_SIZE)
        device_cert = generate_device_cert(device_private_key, ca_private_key, ca_cert, device_id)
        
        device_dir = f"{OUTPUT_DIR}/device_{device_id}"
        save_to_file(device_private_key, f"{device_dir}/device_private_key.pem", is_private_key=True)
        save_to_file(device_cert, f"{device_dir}/device_cert.pem")
        save_to_file(device_private_key.public_key(), f"{device_dir}/device_public_key.pem", is_public_key=True)
        save_to_file(ca_cert, f"{device_dir}/ca_cert.pem")  # 包含CA证书副本，方便设备使用
        
        # 保存设备密钥对
        device_keys[device_id] = (device_private_key, device_private_key.public_key())
        
        # 打印设备证书和密钥为C语言格式并写入日志
        device_cert_c = format_cert_as_c_variable(device_cert, f"device_{device_id}_certificates")
        device_private_key_c = format_key_as_c_variable(device_private_key, f"device_{device_id}_private_key", is_private_key=True)
        device_public_key_c = format_key_as_c_variable(device_private_key.public_key(), f"device_{device_id}_public_key")
        append_to_log(log_path, device_cert_c)
        append_to_log(log_path, device_private_key_c)
        append_to_log(log_path, device_public_key_c)
    
    # 生成共享密钥并对比
    print("\n生成共享密钥并对比...")
    device_1_private_key, device_1_public_key = device_keys["001"]
    device_2_private_key, device_2_public_key = device_keys["002"]
    
    shared_secret_1 = generate_shared_secret(device_1_private_key, device_2_public_key)
    shared_secret_2 = generate_shared_secret(device_2_private_key, device_1_public_key)
    
    if shared_secret_1 == shared_secret_2:
        print("共享密钥匹配!")
    else:
        print("共享密钥不匹配!")
    
    # 打印共享密钥为十六进制
    print("共享密钥 (十六进制):")
    print("设备1私钥 + 设备2公钥:", shared_secret_1.hex())
    print("设备2私钥 + 设备1公钥:", shared_secret_2.hex())
    
    # 保存共享密钥到日志文件
    shared_secret_c = format_shared_secret_as_c_array(shared_secret_1, "shared_secret")
    append_to_log(log_path, shared_secret_c)

    print("\n证书生成完成!")
    print(f"文件已保存到: {os.path.abspath(OUTPUT_DIR)}")
    print(f"日志文件位置: {os.path.abspath(log_path)}")
    print("\nCA证书位置:")
    print(f"  {os.path.abspath(f'{OUTPUT_DIR}/ca/ca_cert.pem')}")

if __name__ == "__main__":
    main()
